
import os
import json
from pydantic import BaseModel, Field
from dashscope import Generation

# 定义你的输出对象
class Date(BaseModel):
    year: int = Field(description="Year")
    month: int = Field(description="Month")
    day: int = Field(description="Day")
    era: str = Field(description="BC or AD")

# 从环境变量获取API密钥
api_key = os.getenv("DASHSCOPE_API_KEY")
if not api_key:
    api_key = "sk-2c036d80d14a443f8769e072bb09a2e1"

def extract_date_with_structured_output(query):
    """使用DashScope提取日期信息并返回结构化输出"""
    
    # 构建提示词，要求返回JSON格式
    prompt = f"""提取用户输入中的日期信息，并以JSON格式返回结果。

要求：
- year: 年份（数字）
- month: 月份（数字，1-12）
- day: 日期（数字，1-31）
- era: 时代（"BC" 或 "AD"）

用户输入: {query}

请只返回JSON格式的结果，不要包含其他文字。"""

    try:
        response = Generation.call(
            model='qwen-plus',
            prompt=prompt,
            api_key=api_key
        )
        
        # 提取响应文本
        response_text = response.output.text.strip()
        
        # 尝试解析JSON
        try:
            # 如果响应包含其他文字，尝试提取JSON部分
            if "```json" in response_text:
                json_start = response_text.find("```json") + 7
                json_end = response_text.find("```", json_start)
                json_text = response_text[json_start:json_end].strip()
            elif "{" in response_text and "}" in response_text:
                json_start = response_text.find("{")
                json_end = response_text.rfind("}") + 1
                json_text = response_text[json_start:json_end]
            else:
                json_text = response_text
            
            # 解析JSON并创建Date对象
            data = json.loads(json_text)
            date_obj = Date(**data)
            return date_obj
            
        except (json.JSONDecodeError, ValueError) as e:
            print(f"JSON解析错误: {e}")
            print(f"原始响应: {response_text}")
            return None
            
    except Exception as e:
        print(f"调用模型时出错: {e}")
        return None

# 测试结构化输出
query = "2023年四月6日天气晴..."
print(f"用户输入: {query}")
print("="*50)

result = extract_date_with_structured_output(query)

if result:
    print("提取的日期信息:")
    print(f"年份: {result.year}")
    print(f"月份: {result.month}")
    print(f"日期: {result.day}")
    print(f"时代: {result.era}")
    print("\n完整对象:")
    print(result)
else:
    print("日期提取失败")
